package net.cyue.ort.llm.data;

import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * 模型输入数据类
 * 封装模型推理所需的输入信息，包括token IDs、缓存等
 * 优化后支持直接指定pastSeqLen，提高代码可维护性和性能
 */
public class ModelInput {
    private final List<Long> inputIds;
    private final Map<String, Object> pastKeyValueCache; // 泛型化缓存类型以支持不同实现
    private final boolean isFirstRun;
    private final int pastSeqLen; // 添加pastSeqLen字段，直接指定历史序列长度
    
    /**
     * 构造函数，不指定pastSeqLen（保持向后兼容）
     * pastSeqLen将从缓存中推断，默认为0
     */
    public ModelInput(List<Long> inputIds, Map<String, Object> pastKeyValueCache, boolean isFirstRun) {
        this(inputIds, pastKeyValueCache, isFirstRun, 0);
    }
    
    /**
     * 构造函数，显式指定pastSeqLen
     * @param inputIds 输入token IDs
     * @param pastKeyValueCache 过去的键值缓存
     * @param isFirstRun 是否是第一次运行
     * @param pastSeqLen 过去的序列长度
     */
    public ModelInput(List<Long> inputIds, Map<String, Object> pastKeyValueCache, boolean isFirstRun, int pastSeqLen) {
        this.inputIds = inputIds != null ? List.copyOf(inputIds) : Collections.emptyList();
        this.pastKeyValueCache = pastKeyValueCache != null ? pastKeyValueCache : Collections.emptyMap();
        this.isFirstRun = isFirstRun;
        this.pastSeqLen = Math.max(0, pastSeqLen); // 确保pastSeqLen非负
    }
    
    /**
     * 创建第一次运行的ModelInput
     * @param inputIds 输入token IDs
     * @return 第一次运行的ModelInput实例
     */
    public static ModelInput createFirstRun(List<Long> inputIds) {
        return new ModelInput(inputIds, Collections.emptyMap(), true, 0);
    }
    
    /**
     * 创建后续运行的ModelInput
     * @param inputIds 输入token IDs
     * @param pastKeyValueCache 过去的键值缓存
     * @param pastSeqLen 过去的序列长度
     * @return 后续运行的ModelInput实例
     */
    public static ModelInput createSubsequentRun(List<Long> inputIds, Map<String, Object> pastKeyValueCache, int pastSeqLen) {
        return new ModelInput(inputIds, pastKeyValueCache, false, pastSeqLen);
    }
    
    public List<Long> getInputIds() {
        return inputIds;
    }
    
    public Map<String, Object> getPastKeyValueCache() {
        return pastKeyValueCache;
    }
    
    public boolean isFirstRun() {
        return isFirstRun;
    }
    
    /**
     * 获取过去的序列长度
     * @return 过去的序列长度
     */
    public int getPastSeqLen() {
        return pastSeqLen;
    }
    
    /**
     * 检查是否有有效的pastSeqLen
     * @return 如果pastSeqLen大于0返回true，否则返回false
     */
    public boolean hasValidPastSeqLen() {
        return pastSeqLen > 0;
    }
    
    @Override
    public String toString() {
        return "ModelInput{" +
            "inputIds=" + inputIds +
            ", isFirstRun=" + isFirstRun +
            ", pastSeqLen=" + pastSeqLen +
            ", hasCache=" + !pastKeyValueCache.isEmpty() +
            '}';
    }
    
    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (o == null || getClass() != o.getClass()) return false;
        ModelInput that = (ModelInput) o;
        return isFirstRun == that.isFirstRun &&
            pastSeqLen == that.pastSeqLen &&
            Objects.equals(inputIds, that.inputIds) &&
            Objects.equals(pastKeyValueCache, that.pastKeyValueCache);
    }
    
    @Override
    public int hashCode() {
        return Objects.hash(inputIds, pastKeyValueCache, isFirstRun, pastSeqLen);
    }
}

